import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
from nicegui import ui
import numpy as np


def create_survival_by_class(df):
    """三维：舱位-生存状态-人数（美观优化）"""
    try:
        df = df.copy()
        df['Pclass'] = df['Pclass'].astype(int)
        df['Survived'] = df['Survived'].astype(int)
        survival_df = df.groupby(['Pclass', 'Survived']).size().reset_index(name='Count')
        survival_df['SurvivedLabel'] = survival_df['Survived'].map({0: '遇难', 1: '幸存'})
        fig = px.scatter_3d(
            survival_df,
            x='Pclass', y='Survived', z='Count',
            color='SurvivedLabel',
            title='不同舱位等级与生存状态的三维分布',
            labels={'Pclass': '舱位等级', 'Survived': '生存状态', 'Count': '人数'},
            color_discrete_map={'遇难': '#EF553B', '幸存': '#00CC96'},
            hover_data=['Pclass', 'SurvivedLabel', 'Count'],
        )
        fig.update_traces(marker=dict(size=4, opacity=0.85, line=dict(width=2, color='DarkSlateGrey')))
        fig.update_layout(
            height=500,
            margin=dict(l=80, r=80, t=100, b=60),
            showlegend=True,
            legend=dict(font=dict(size=16)),
            scene=dict(
                xaxis_title='舱位等级',
                yaxis_title='生存状态',
                zaxis_title='人数',
                xaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                yaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                zaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                bgcolor='white',
                camera=dict(eye=dict(x=1.5, y=1.5, z=1.2))
            )
        )
        return fig
    except Exception as e:
        ui.notify(f"创建三维舱位生存率图表失败: {str(e)}", type='negative')
        return None


def create_age_distribution(df):
    """三维：年龄-生存状态-人数（美观优化）"""
    try:
        df = df.copy()
        df['Age'] = pd.to_numeric(df['Age'], errors='coerce')
        df['Survived'] = df['Survived'].astype(int)
        df_age = df[df['Age'].notna()]
        # 分箱统计
        df_age['AgeBin'] = pd.cut(df_age['Age'], bins=10)
        age_group = df_age.groupby(['AgeBin', 'Survived']).size().reset_index(name='Count')
        age_group['AgeMid'] = age_group['AgeBin'].apply(lambda x: x.mid)
        age_group['SurvivedLabel'] = age_group['Survived'].map({0: '遇难', 1: '幸存'})
        fig = px.scatter_3d(
            age_group,
            x='AgeMid', y='Survived', z='Count',
            color='SurvivedLabel',
            title='年龄-生存状态-人数三维分布',
            labels={'AgeMid': '年龄', 'Survived': '生存状态', 'Count': '人数'},
            color_discrete_map={'遇难': '#EF553B', '幸存': '#00CC96'},
            hover_data=['AgeMid', 'SurvivedLabel', 'Count'],
        )
        fig.update_traces(marker=dict(size=4, opacity=0.85, line=dict(width=2, color='DarkSlateGrey')))
        fig.update_layout(
            height=500,
            margin=dict(l=80, r=80, t=100, b=60),
            showlegend=True,
            legend=dict(font=dict(size=16)),
            scene=dict(
                xaxis_title='年龄',
                yaxis_title='生存状态',
                zaxis_title='人数',
                xaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                yaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                zaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
                bgcolor='white',
                camera=dict(eye=dict(x=1.5, y=1.5, z=1.2))
            )
        )
        return fig
    except Exception as e:
        ui.notify(f"创建三维年龄分布图表失败: {str(e)}", type='negative')
        return None


def create_fare_vs_age(df):
    """二维：年龄-票价（按生存状态）"""
    try:
        df = df.copy()
        df['Age'] = pd.to_numeric(df['Age'], errors='coerce')
        df['Fare'] = pd.to_numeric(df['Fare'], errors='coerce')
        df['Survived'] = df['Survived'].astype(int)
        df_clean = df[df['Age'].notna() & df['Fare'].notna()]
        fig = px.scatter(
            df_clean,
            x='Age', y='Fare',
            color='Survived',
            title='票价-年龄二维关系（按生存状态）',
            labels={'Age': '年龄', 'Fare': '票价'},
            color_discrete_map={0: '#EF553B', 1: '#00CC96'},
            hover_data=['Age', 'Fare', 'Pclass', 'Survived'],
        )
        fig.update_traces(marker=dict(size=6, opacity=0.7, line=dict(width=0)))
        fig.update_layout(
            height=500,
            margin=dict(l=80, r=80, t=100, b=60),
            showlegend=True,
            legend=dict(font=dict(size=16)),
            xaxis_title='年龄',
            yaxis_title='票价',
            xaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
            yaxis=dict(title_font=dict(size=18), tickfont=dict(size=14), gridcolor='lightgrey'),
            plot_bgcolor='white',
            paper_bgcolor='white',
        )
        return fig
    except Exception as e:
        ui.notify(f"创建票价年龄二维关系图表失败: {str(e)}", type='negative')
        return None


def create_correlation_heatmap(df):
    """二维热力图保持不变"""
    try:
        df = df.copy()
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        df_numeric = df[numeric_columns]
        corr = df_numeric.corr().round(2)
        fig = px.imshow(
            corr,
            text_auto=True,
            title='特征相关性热力图',
            color_continuous_scale='RdBu_r',
            zmin=-1,
            zmax=1,
            aspect='auto'
        )
        fig.update_layout(
            height=500,
            margin=dict(l=100, r=100, t=100, b=60),
            showlegend=True,
            hovermode='closest',
            plot_bgcolor='white',
            paper_bgcolor='white',
            autosize=True
        )
        return fig
    except Exception as e:
        ui.notify(f"创建相关性热力图失败: {str(e)}", type='negative')
        return None